from typing import Optional, List

from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel

from .base import BaseNodes, tool_node
from ...models.graph import GraphState


class AskUserQuestionSchema(BaseModel):
    question: str


class AskUserSelectSchema(BaseModel):
    options: List[str]
    question: str
    multiple: Optional[bool] = False


class AskUserNodes(BaseNodes):

    @tool_node(args_schema=AskUserQuestionSchema)
    def ask_user_question(self, state: GraphState, *, config: Optional[RunnableConfig] = None):
        """
        Description: Ask the user a question to gather additional information needed to complete the task. This tool should be used when you encounter ambiguities, need clarification, or require more details to proceed effectively. It allows for interactive problem-solving by enabling direct communication with the user. Use this tool judiciously to maintain a balance between gathering necessary information and avoiding excessive back-and-forth.
        Parameters:
        - question: (required) The question to ask the user. This should be a clear, specific question that addresses the information you need.
        Usage:
        <ask_user_question>
        <question>Your question here</question>
        </ask_user_question>
        """
        return self.wait_human_feedback()

    @tool_node(args_schema=AskUserSelectSchema)
    def ask_user_select(self, state: GraphState, *, config: Optional[RunnableConfig] = None):
        """
        Description: Present a list of options to the user and ask them to make a selection. This tool should be used when you need the user to choose from predefined options. By default it allows single selection, but can be configured for multiple selections. Use this when you need to guide user decisions through specific choices rather than open-ended questions.
        Parameters:
        - options: (required) A list of options for the user to choose from, with one option per line. Each option should be clear and distinct.
        - question: (required) The question or prompt to present to the user along with the options.
        - multiple: (optional) Boolean flag to indicate whether multiple selections are allowed. Defaults to false (single selection).
        Usage:
        <ask_user_select>
        <options>
        Option 1
        Option 2
        Option 3
        </options>
        <question>Your question here</question>
        <multiple>true or false (optional, defaults to false)</multiple>
        </ask_user_select>
        """
        return self.wait_human_feedback()
